import cv2
import numpy as np
from ultralytics import YOLO

# Load the model
model = YOLO("yolo11n-seg.pt")

# Path to the image
img_path = '/home/JSDC/017254/code/gitee/map_learing/map_grid/04_xyzrpy_visual/demo/bus.jpg'

# Predict with the model
results = model(img_path)

# Load the original image to get dimensions
image = cv2.imread(img_path)
height, width = image.shape[:2]

# Iterate over the results and save each mask individually
mask_index = 0  # Initialize mask index
for result in results:
    # Access the mask data
    masks = result.masks.data

    # Iterate over each mask
    for mask in masks:
        # Convert the mask to a NumPy array and then to 0-255 format
        mask = mask.cpu().numpy().astype(np.uint8) * 255  # Convert Tensor to NumPy array, then to 0-255 format

        # Resize mask to match the original image dimensions if needed
        mask = cv2.resize(mask, (width, height))
        print(f'mask.shape {mask.shape}')
        print(f'mask.sum {mask.sum() / 255}')

        # Create a unique filename for each mask
        filename = f"mask_{mask_index}.png"

        # Save the mask as a black and white image
        cv2.imwrite(filename, mask)
        print(f"Saved mask as {filename}")

        # Increment the mask index
        mask_index += 1
